import os
import csv

def read_transcript(file_dir=None, file_name="aishell_transcript_v0.8.txt"):
    txt_dict = {} # file_name: txt
    line_list = []
    if file_dir is not None:
        file_name = os.path.join(file_dir, file_name)
    with open(file_name,"r", encoding="utf-8-sig") as txt_file:
        line_list = txt_file.readlines()
    for i, line in enumerate(line_list):
        line_strip = line.strip()
        txt_dict[line_strip[:16]] = line_strip[16:].strip()
    return txt_dict

def read_csv(file_dir=None, file_name="result.csv"):
    csv_dict = {}
    if file_dir is not None:
        file_name = os.path.join(file_dir, file_name)
    print("Calculate on file result: %s"%(file_name))
    with open(file_name, "r", encoding="utf-8-sig") as csv_file:
        csv_content = csv.reader(csv_file)
        # header = next(csv_content)
        for row in csv_content:
            csv_dict[row[0]] = row[1]
    return csv_dict

def get_csv_txt(csv_dict, txt_dict):
    merge = {}
    for k, v in csv_dict.items():
        key = k.split(".")[0]
        if key in txt_dict.keys():
            merge[key] = [v, txt_dict[key]]
    return merge

if __name__ == "__main__":
    file_dir = "/public/ai_platform/models/cmcc_2022/dataset_aishell/transcript"
    txt_dict = read_transcript(file_dir=file_dir)

    csv_dir = "/home/gyf/pkg/xxgg/github/ai_app/cmcc/asr_conformer/tensorrt"
    csv_file = "result_test.csv"
    csv_dict = read_csv(csv_dir, csv_file)

    # print(csv_dict)
    # print(txt_dict)
    get_csv_txt(csv_dict, txt_dict)